import json
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path


def get_mean_at_step(results, step):
    results_at_a_step = [result[step] for result in results]
    return np.mean(results_at_a_step)


def get_std_at_step(results, step):
    results_at_a_step = [result[step] for result in results]
    return np.std(results_at_a_step)


def get_mean_and_std_differ_results(results):
    mean_results = [get_mean_at_step(results, step) for step in range(len(results[0]))]
    mean_minus_std_results = [mean - get_std_at_step(results, step) for step, mean in enumerate(mean_results)]
    mean_plus_std_results = [mean + get_std_at_step(results, step) for step, mean in enumerate(mean_results)]
    return mean_results, mean_minus_std_results, mean_plus_std_results


agent_to_color_dict = {
    "DQN": "#E50000",  # red
    "D3QN": "#9A0EEA",  # violet
    "DDPG": "#008000",# green
    "TD3": "#0343DF",# blue
    'PPODiscrete': '#FF00FF', # pink
    'PPOContinuous': '#FF00FF' # pink
}

if __name__ == '__main__':
    env_name = 'Humanoid-v5'
    path = Path(f'{env_name}.json')
    contents = path.read_text()
    contents = json.loads(contents)
    plt.style.use('seaborn')
    fig, ax = plt.subplots()
    for agent_name in contents.keys():
        color = agent_to_color_dict[agent_name]
        agent_results = contents[agent_name]
        achieved_goal = agent_results[0][2]
        results = [result[0] for result in agent_results] # idx 0:eval 1:achieved_goal 2:seed
        avg_results = [[] for _ in range(len(results))]
        for i, result in enumerate(results):
            for j in range(len(result)):
                l_ind = max(j - 20, 0) #j
                r_ind = j + 1 #min(j + 30, len(result))
                avg_results[i].append(np.mean(result[l_ind:r_ind]))
        means, mean_minus_stds, mean_plus_stds = get_mean_and_std_differ_results(avg_results)
        x = list(range(len(means)))

        ax.set_title(f'{env_name}', fontsize=15)
        ax.set_xlabel('Episodes', fontsize=14)
        ax.set_ylabel('Rolling Episode Scores', fontsize=14)

        ax.plot(x, means, label=agent_name, color=color, linewidth=1.3)
        ax.plot(x, mean_minus_stds, alpha=0.1, color=color, linewidth=1.3)
        ax.plot(x, mean_plus_stds, alpha=0.1, color=color, linewidth=1.3)
        ax.fill_between(x, y1=mean_minus_stds, y2=mean_plus_stds, alpha=0.1, color=color)

        legend = ax.legend(loc='center', bbox_to_anchor=(0.5, -0.3), shadow=True, fancybox=True, ncol=3, fontsize='large')
        # legend = ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.15), fancybox=True, shadow=True, ncol=3)
        legend.get_frame().set_facecolor('xkcd:white')
        ax.set_xlim([0, x[-1] + 1])
    # plt.show()
    plt.savefig(f'{env_name}.png', bbox_inches='tight')